#! /usr/bin/env python
# coding=utf-8
# infile = r'test123.fq'

# outfile = r'result111'



from Bio import SeqIO
import os
import shutil
import argparse
import sys
import re

# 使用方法 
# conver04_add_gap2simulation.py  -r rawfile -s simulationfile 

parser = argparse.ArgumentParser(
    description=''' 将minibar.py生成的 分配好sample id 的fastq 文件 按每个 amplicon 一个文件进行提取 并自动翻转
    用法:
    barcode_assign_to_file.py  -i result.fq -o result
    由大天才于2022年2月17日创建于浙江农业大学''')

parser.add_argument('-i',
                help='minibar.py输出的分配好的fastq文件')

parser.add_argument('-o',
                help='输出的目录')

args = parser.parse_args()


if not args.i or not args.o:
    parser.print_help()
    sys.exit()




infile = args.i

outfile = args.o



try:
    shutil.rmtree(outfile)
except:
    pass
try:
    os.mkdir(outfile)
except:
    pass



match_reverse = re.compile(r'H[-].+H[+].+')
match_item = re.compile(r'H[-|+].+H[-|+].+')



k = 0
for i in SeqIO.parse(infile,'fastq'):
    k += 1
    if k%10000==0:
        print(k)

    if len(i.description)>999:
        continue
    else:
        d = match_item.findall(i.description)

        if d!=[] and len(d[0].split(' '))==2:
            
            group = d[0].split(' ')[1]

            if re.findall(match_reverse,str(i.description))!=[]:
                tmp_id= i.id
                tmp_des = i.description
                i = i.reverse_complement()
                i.id= tmp_id
                i.description = tmp_des

            if os.path.exists(outfile+'/'+group+'.fastq'):
                h = open(outfile+'/'+group+'.fastq','a')
                SeqIO.write(i,h,'fastq')
                h.close()
            else:
                h = open(outfile+'/'+group+'.fastq','w')
                SeqIO.write(i,h,'fastq')
                h.close()
        